#!/usr/bin/env python3
"""
Phase 4f: Two analyses
1. Horse race: Z×KAOPEN vs Z×log(GDP_pc) in developing subsample
   - Is the interaction really about financial openness, or just income/development?
2. Leave-one-region-out jackknife on the baseline model
   - Are the larger 140c coefficients stable, or driven by specific country groups?
"""

import sys
import pandas as pd
import numpy as np
from pathlib import Path
from scipy import stats

FOLLOWUP_DIR = Path("/mnt/c/demographics_capital_flows/multilateral/followup")
PROJECT_DIR = FOLLOWUP_DIR.parent
sys.path.insert(0, str(FOLLOWUP_DIR))
sys.path.insert(1, str(PROJECT_DIR))

from src.model import PanelGLS
from src.macro import EBA_COUNTRIES, SSA_COUNTRIES, EU_EXPANSION, EXPANSION_TIER1, filter_eba_sample

OUTPUT_DIR = FOLLOWUP_DIR / "output" / "tables"
PROCESSED_DIR = FOLLOWUP_DIR / "data" / "processed"

panel = pd.read_csv(PROCESSED_DIR / "full_panel.csv")

est = panel[
    (panel['ca_gdp'].notna()) &
    (panel['year'] >= 1986) &
    (panel['year'] <= 2024)
].copy()

full_sample = filter_eba_sample(est, extended=True, expansion=True)
print(f"Full sample: {full_sample['iso3'].nunique()} countries, {len(full_sample)} obs")

demo_vars = ['Z_1', 'Z_2', 'Z_3']
baseline_controls = ['fiscal_bal_gdp', 'kaopen', 'expected_growth', 'nfa_gdp_lag',
                     'log_rel_opw', 'health_exp_gdp']


def run_gls(df, dep_var, indep_vars):
    """Run PanelGLS and return a dict of results keyed by variable name."""
    comp = df.dropna(subset=[dep_var] + indep_vars).copy()
    y = comp[dep_var].values
    X = comp[indep_vars].values
    gls = PanelGLS()
    gls.fit(y, X, comp['iso3'].values, comp['year'].values)
    result = {
        'r_squared': gls.r_squared,
        'n_obs': gls.n_obs,
        'n_countries': gls.n_countries,
        'coefficients': dict(zip(indep_vars, gls.beta)),
        'std_errors': dict(zip(indep_vars, gls.se)),
        'p_values': dict(zip(indep_vars, gls.pvalues)),
    }
    return result


# ═══════════════════════════════════════════════════════════════════════════
# 1. HORSE RACE: Z×KAOPEN vs Z×log(GDP_pc) in developing subsample
# ═══════════════════════════════════════════════════════════════════════════
print("=" * 70)
print("HORSE RACE: Z×KAOPEN vs Z×log(GDP_pc)")
print("=" * 70)

# Income classification
if 'gdp_pc_ppp' in full_sample.columns:
    gdp_pc = full_sample.groupby('iso3')['gdp_pc_ppp'].median()
else:
    full_sample['gdp_pc_approx'] = full_sample['ngdp_usd'] / full_sample['population_weo']
    gdp_pc = full_sample.groupby('iso3')['gdp_pc_approx'].median()

q67 = gdp_pc.quantile(0.67)
income_map = {iso: ('high' if val >= q67 else 'low_mid') for iso, val in gdp_pc.items()}
for iso, val in gdp_pc.items():
    if pd.isna(val):
        income_map[iso] = 'low_mid'

full_sample['income_group'] = full_sample['iso3'].map(income_map)

# Create log GDP per capita (demeaned for interpretability)
gdp_pc_series = gdp_pc.reindex(full_sample['iso3'].values)
full_sample['log_gdp_pc'] = np.log(gdp_pc_series.values + 1)
full_sample['log_gdp_pc_dm'] = full_sample['log_gdp_pc'] - full_sample['log_gdp_pc'].mean()

# Create interactions
for z in demo_vars:
    full_sample[f'{z}_x_kaopen'] = full_sample[z] * full_sample['kaopen']
    full_sample[f'{z}_x_log_gdp'] = full_sample[z] * full_sample['log_gdp_pc_dm']

kaopen_ints = [f'{z}_x_kaopen' for z in demo_vars]
gdp_ints = [f'{z}_x_log_gdp' for z in demo_vars]

# Developing subsample (low + middle income)
dev_sample = full_sample[full_sample['income_group'] == 'low_mid'].copy()
print(f"\nDeveloping subsample: {dev_sample['iso3'].nunique()} countries")

# Use baseline controls (no lending rate) for larger sample
all_hr_vars = demo_vars + baseline_controls + kaopen_ints + gdp_ints
dev_base = dev_sample.dropna(subset=all_hr_vars + ['ca_gdp']).copy()
print(f"Estimation sample: {dev_base['iso3'].nunique()} countries, {len(dev_base)} obs")

# Correlation check
corr = dev_base['kaopen'].corr(dev_base['log_gdp_pc'])
print(f"Correlation(KAOPEN, log GDP/pc) in developing sample: {corr:.3f}")

# Model H0: Baseline (no interactions)
r_base = run_gls(dev_base, 'ca_gdp', demo_vars + baseline_controls)
print(f"\n  H0 (baseline):        R² = {r_base['r_squared']:.4f}")

# Model H1: + Z×KAOPEN only
r_kaopen = run_gls(dev_base, 'ca_gdp', demo_vars + baseline_controls + kaopen_ints)
print(f"  H1 (+ Z×KAOPEN):     R² = {r_kaopen['r_squared']:.4f}")

# Model H2: + Z×log(GDP_pc) only
r_gdp = run_gls(dev_base, 'ca_gdp', demo_vars + baseline_controls + gdp_ints)
print(f"  H2 (+ Z×log GDP):    R² = {r_gdp['r_squared']:.4f}")

# Model H3: Both (horse race)
r_both = run_gls(dev_base, 'ca_gdp', demo_vars + baseline_controls + kaopen_ints + gdp_ints)
print(f"  H3 (both):            R² = {r_both['r_squared']:.4f}")

# Report coefficients
for label, result, vars_of_interest in [
    ('H1: Z×KAOPEN only', r_kaopen, kaopen_ints),
    ('H2: Z×log(GDP_pc) only', r_gdp, gdp_ints),
    ('H3: Horse race (both)', r_both, kaopen_ints + gdp_ints),
]:
    print(f"\n  --- {label} ---")
    for v in vars_of_interest:
        c = result['coefficients'].get(v, np.nan)
        p = result['p_values'].get(v, np.nan)
        sig = '***' if p < 0.001 else '**' if p < 0.01 else '*' if p < 0.05 else ''
        print(f"    {v:20s}: {c:8.3f} (p={p:.4f}) {sig}")

# Joint F-tests in horse race
n_hr = r_both['n_obs']
k_hr = len(demo_vars + baseline_controls + kaopen_ints + gdp_ints)

# F-test: are KAOPEN interactions jointly zero in H3?
q = 3
F_kaopen = ((r_both['r_squared'] - r_gdp['r_squared']) / q) / ((1 - r_both['r_squared']) / (n_hr - k_hr - 1))
p_kaopen_f = 1 - stats.f.cdf(F_kaopen, q, n_hr - k_hr - 1)

# F-test: are GDP interactions jointly zero in H3?
F_gdp = ((r_both['r_squared'] - r_kaopen['r_squared']) / q) / ((1 - r_both['r_squared']) / (n_hr - k_hr - 1))
p_gdp_f = 1 - stats.f.cdf(F_gdp, q, n_hr - k_hr - 1)

print(f"\n  Joint F-tests in horse race (H3):")
print(f"    KAOPEN interactions jointly zero? F({q},{n_hr-k_hr-1}) = {F_kaopen:.3f}, p = {p_kaopen_f:.4f}")
print(f"    GDP interactions jointly zero?    F({q},{n_hr-k_hr-1}) = {F_gdp:.3f}, p = {p_gdp_f:.4f}")

if p_kaopen_f < 0.05 and p_gdp_f >= 0.05:
    print("    >>> KAOPEN WINS: openness matters beyond income level")
elif p_gdp_f < 0.05 and p_kaopen_f >= 0.05:
    print("    >>> GDP WINS: income level matters beyond openness")
elif p_kaopen_f < 0.05 and p_gdp_f < 0.05:
    print("    >>> BOTH SURVIVE: independent channels")
else:
    print("    >>> NEITHER SURVIVES: multicollinearity or weak effects")

# Save horse race results
hr_rows = []
for model_name, result in [
    ('H0_baseline', r_base),
    ('H1_kaopen_only', r_kaopen),
    ('H2_gdp_only', r_gdp),
    ('H3_horse_race', r_both),
]:
    for v in result['coefficients']:
        hr_rows.append({
            'model': model_name,
            'variable': v,
            'coefficient': result['coefficients'][v],
            'std_error': result['std_errors'][v],
            'p_value': result['p_values'][v],
            'r_squared': result['r_squared'],
            'n_obs': result['n_obs'],
        })

hr_df = pd.DataFrame(hr_rows)
hr_df.to_csv(OUTPUT_DIR / 'kaopen_vs_gdp_horserace.csv', index=False)
print(f"\n  Saved to {OUTPUT_DIR / 'kaopen_vs_gdp_horserace.csv'}")


# ═══════════════════════════════════════════════════════════════════════════
# 2. LEAVE-ONE-REGION-OUT JACKKNIFE
# ═══════════════════════════════════════════════════════════════════════════
print("\n" + "=" * 70)
print("LEAVE-ONE-REGION-OUT JACKKNIFE: BASELINE MODEL")
print("=" * 70)

# Region classification
all_countries = full_sample['iso3'].unique()

latin_am = {'ARG', 'BOL', 'BRA', 'CHL', 'COL', 'CRI', 'DOM', 'ECU', 'GTM', 'HND',
            'MEX', 'NIC', 'PAN', 'PER', 'PRY', 'SLV', 'URY', 'VEN'}
mena = {'DZA', 'BHR', 'EGY', 'IRN', 'IRQ', 'ISR', 'JOR', 'KWT', 'LBN', 'LBY',
        'MAR', 'OMN', 'QAT', 'SAU', 'TUN', 'ARE', 'YEM', 'PSE', 'SYR'}
cca = {'ARM', 'AZE', 'BLR', 'GEO', 'KAZ', 'KGZ', 'MDA', 'MNG', 'RUS', 'TJK',
       'TKM', 'UKR', 'UZB'}
south_se_asia = {'BGD', 'BTN', 'IND', 'IDN', 'KHM', 'LAO', 'LKA', 'MMR', 'MYS',
                 'NPL', 'PAK', 'PHL', 'SGP', 'THA', 'VNM'}
east_asia = {'CHN', 'HKG', 'JPN', 'KOR', 'TWN'}
adv_europe = {'AUT', 'BEL', 'CHE', 'CYP', 'CZE', 'DEU', 'DNK', 'ESP', 'EST', 'FIN',
              'FRA', 'GBR', 'GRC', 'HRV', 'HUN', 'IRL', 'ISL', 'ITA', 'LTU', 'LUX',
              'LVA', 'MLT', 'NLD', 'NOR', 'POL', 'PRT', 'ROU', 'BGR', 'SVK', 'SVN',
              'SWE'}
adv_other = {'AUS', 'CAN', 'NZL', 'USA'}
ssa_un = {'AGO', 'BDI', 'BEN', 'BFA', 'BWA', 'CAF', 'CIV', 'CMR', 'COD', 'COG',
          'COM', 'CPV', 'DJI', 'ERI', 'ETH', 'GAB', 'GHA', 'GIN', 'GMB', 'GNB',
          'GNQ', 'KEN', 'LBR', 'LSO', 'MDG', 'MLI', 'MOZ', 'MRT', 'MUS', 'MWI',
          'NAM', 'NER', 'NGA', 'RWA', 'SDN', 'SEN', 'SLE', 'SOM', 'SSD', 'STP',
          'SWZ', 'SYC', 'TCD', 'TGO', 'TZA', 'UGA', 'ZAF', 'ZMB', 'ZWE'}

region_map = {}
for c in all_countries:
    if c in ssa_un or c in set(SSA_COUNTRIES):
        region_map[c] = 'SSA'
    elif c in mena:
        region_map[c] = 'MENA'
    elif c in cca:
        region_map[c] = 'CCA'
    elif c in latin_am:
        region_map[c] = 'LatAm'
    elif c in south_se_asia:
        region_map[c] = 'S/SE_Asia'
    elif c in east_asia:
        region_map[c] = 'E_Asia'
    elif c in adv_europe:
        region_map[c] = 'Adv_Europe'
    elif c in adv_other:
        region_map[c] = 'Adv_Other'
    else:
        region_map[c] = 'Other'

full_sample['region'] = full_sample['iso3'].map(region_map)

# Show region composition
print("\nRegion composition:")
region_counts = full_sample.groupby('region')['iso3'].nunique().sort_values(ascending=False)
for reg, n in region_counts.items():
    countries = sorted(full_sample[full_sample['region'] == reg]['iso3'].unique())
    print(f"  {reg:15s}: {n:3d} countries  ({', '.join(countries[:8])}{'...' if len(countries) > 8 else ''})")

# Baseline estimation on full sample
base_vars = demo_vars + baseline_controls
base_sample = full_sample.dropna(subset=base_vars + ['ca_gdp']).copy()
base_sample['region'] = base_sample['iso3'].map(region_map)
print(f"\nBaseline estimation sample: {base_sample['iso3'].nunique()} countries, {len(base_sample)} obs")

r_full = run_gls(base_sample, 'ca_gdp', base_vars)

print(f"\n  FULL SAMPLE: R² = {r_full['r_squared']:.4f}")
for v in demo_vars:
    c = r_full['coefficients'][v]
    p = r_full['p_values'][v]
    print(f"    {v}: {c:8.3f} (p={p:.4f})")

# Jackknife: drop one region at a time
jackknife_results = []
regions_to_drop = sorted(base_sample['region'].unique())

print(f"\n--- Leave-one-region-out ---")
print(f"  {'Dropped':<15s} {'N_c':>4s} {'N_obs':>6s} {'R²':>7s} {'Z₁':>9s} {'p(Z₁)':>8s} {'Z₂':>9s} {'p(Z₂)':>8s} {'Z₃':>9s} {'p(Z₃)':>8s}")
print("  " + "-" * 93)

for region in regions_to_drop:
    dropped_countries = base_sample[base_sample['region'] == region]['iso3'].unique()
    sub = base_sample[base_sample['region'] != region].copy()
    n_c = sub['iso3'].nunique()

    r_jack = run_gls(sub, 'ca_gdp', base_vars)

    row = {
        'dropped_region': region,
        'n_countries': n_c,
        'n_countries_dropped': len(dropped_countries),
        'n_obs': r_jack['n_obs'],
        'r_squared': r_jack['r_squared'],
    }
    for v in base_vars:
        row[f'{v}_coef'] = r_jack['coefficients'].get(v, np.nan)
        row[f'{v}_pval'] = r_jack['p_values'].get(v, np.nan)

    jackknife_results.append(row)

    z1c = r_jack['coefficients']['Z_1']
    z1p = r_jack['p_values']['Z_1']
    z2c = r_jack['coefficients']['Z_2']
    z2p = r_jack['p_values']['Z_2']
    z3c = r_jack['coefficients']['Z_3']
    z3p = r_jack['p_values']['Z_3']

    print(f"  {region:<15s} {n_c:4d} {r_jack['n_obs']:6d} {r_jack['r_squared']:7.4f} {z1c:9.2f} {z1p:8.4f} {z2c:9.3f} {z2p:8.4f} {z3c:9.4f} {z3p:8.4f}")

# Add full sample row
print(f"  {'FULL':<15s} {r_full['n_countries']:4d} {r_full['n_obs']:6d} {r_full['r_squared']:7.4f} {r_full['coefficients']['Z_1']:9.2f} {r_full['p_values']['Z_1']:8.4f} {r_full['coefficients']['Z_2']:9.3f} {r_full['p_values']['Z_2']:8.4f} {r_full['coefficients']['Z_3']:9.4f} {r_full['p_values']['Z_3']:8.4f}")

# Coefficient stability summary
jk_df = pd.DataFrame(jackknife_results)

print(f"\n--- Coefficient Stability Summary ---")
for v in demo_vars:
    full_c = r_full['coefficients'][v]
    full_p = r_full['p_values'][v]
    jk_coefs = jk_df[f'{v}_coef']
    min_c = jk_coefs.min()
    max_c = jk_coefs.max()
    range_c = max_c - min_c
    cv = jk_coefs.std() / abs(jk_coefs.mean()) * 100

    # Note: dropping a region that inflates the coefficient means that region
    # was DAMPENING it. Min coefficient = when we drop the amplifying region.
    min_reg = jk_df.loc[jk_coefs.idxmin(), 'dropped_region']
    max_reg = jk_df.loc[jk_coefs.idxmax(), 'dropped_region']

    print(f"\n  {v}:")
    print(f"    Full sample:  {full_c:.3f}")
    print(f"    Range:        [{min_c:.3f}, {max_c:.3f}] (span = {range_c:.3f})")
    print(f"    CV:           {cv:.1f}%")
    print(f"    Smallest when dropping: {min_reg} → {min_c:.3f}")
    print(f"    Largest when dropping:  {max_reg} → {max_c:.3f}")

    sig_changes = []
    for _, row in jk_df.iterrows():
        jp = row[f'{v}_pval']
        if (full_p < 0.05 and jp >= 0.05) or (full_p >= 0.05 and jp < 0.05):
            sig_changes.append(row['dropped_region'])
    if sig_changes:
        print(f"    ⚠ Significance flips when dropping: {', '.join(sig_changes)}")
    else:
        print(f"    Significance stable across all drops")


# ═══════════════════════════════════════════════════════════════════════════
# 3. LEAVE-ONE-REGION-OUT ON EXTENDED (INTERACTION) MODEL
# ═══════════════════════════════════════════════════════════════════════════
print(f"\n\n--- Leave-one-region-out: EXTENDED MODEL (Z×KAOPEN) ---")

for z in demo_vars:
    full_sample[f'{z}_x_kaopen'] = full_sample[z] * full_sample['kaopen']

ext_vars = demo_vars + baseline_controls + ['log_lending_rate'] + [f'{z}_x_kaopen' for z in demo_vars]
ext_sample = full_sample.dropna(subset=ext_vars + ['ca_gdp']).copy()
ext_sample['region'] = ext_sample['iso3'].map(region_map)

print(f"Extended estimation sample: {ext_sample['iso3'].nunique()} countries, {len(ext_sample)} obs")

r_ext_full = run_gls(ext_sample, 'ca_gdp', ext_vars)

print(f"\n  FULL SAMPLE: R² = {r_ext_full['r_squared']:.4f}")
for v in [f'{z}_x_kaopen' for z in demo_vars]:
    c = r_ext_full['coefficients'][v]
    p = r_ext_full['p_values'][v]
    sig = '***' if p < 0.001 else '**' if p < 0.01 else '*' if p < 0.05 else ''
    print(f"    {v}: {c:8.3f} (p={p:.4f}) {sig}")

ext_jk_results = []
print(f"\n  {'Dropped':<15s} {'N_c':>4s} {'N_obs':>6s} {'R²':>7s} {'Z₁×K':>9s} {'p':>7s} {'Z₂×K':>9s} {'p':>7s} {'Z₃×K':>9s} {'p':>7s}")
print("  " + "-" * 88)

for region in sorted(ext_sample['region'].unique()):
    sub = ext_sample[ext_sample['region'] != region].copy()
    n_c = sub['iso3'].nunique()
    if n_c < 10:
        continue

    r_j = run_gls(sub, 'ca_gdp', ext_vars)

    row = {
        'dropped_region': region,
        'n_countries': n_c,
        'n_obs': r_j['n_obs'],
        'r_squared': r_j['r_squared'],
    }
    for v in ext_vars:
        row[f'{v}_coef'] = r_j['coefficients'].get(v, np.nan)
        row[f'{v}_pval'] = r_j['p_values'].get(v, np.nan)
    ext_jk_results.append(row)

    ints = [f'{z}_x_kaopen' for z in demo_vars]
    vals = [(r_j['coefficients'][v], r_j['p_values'][v]) for v in ints]
    print(f"  {region:<15s} {n_c:4d} {r_j['n_obs']:6d} {r_j['r_squared']:7.4f} {vals[0][0]:9.2f} {vals[0][1]:7.4f} {vals[1][0]:9.3f} {vals[1][1]:7.4f} {vals[2][0]:9.4f} {vals[2][1]:7.4f}")

ints = [f'{z}_x_kaopen' for z in demo_vars]
vals_full = [(r_ext_full['coefficients'][v], r_ext_full['p_values'][v]) for v in ints]
print(f"  {'FULL':<15s} {r_ext_full['n_countries']:4d} {r_ext_full['n_obs']:6d} {r_ext_full['r_squared']:7.4f} {vals_full[0][0]:9.2f} {vals_full[0][1]:7.4f} {vals_full[1][0]:9.3f} {vals_full[1][1]:7.4f} {vals_full[2][0]:9.4f} {vals_full[2][1]:7.4f}")

# Interaction stability
ext_jk_df = pd.DataFrame(ext_jk_results)
print(f"\n--- Interaction Stability ---")
for v in ints:
    full_c = r_ext_full['coefficients'][v]
    full_p = r_ext_full['p_values'][v]
    jk_coefs = ext_jk_df[f'{v}_coef']
    jk_pvals = ext_jk_df[f'{v}_pval']
    sig_changes = ext_jk_df[((full_p < 0.05) & (jk_pvals >= 0.05)) | ((full_p >= 0.05) & (jk_pvals < 0.05))]['dropped_region'].tolist()
    print(f"  {v}: full={full_c:.3f} (p={full_p:.4f}), range=[{jk_coefs.min():.3f}, {jk_coefs.max():.3f}]")
    if sig_changes:
        print(f"    ⚠ Significance changes when dropping: {', '.join(sig_changes)}")


# ═══════════════════════════════════════════════════════════════════════════
# 4. CUMULATIVE ADDITION: Build up from EBA-49 → 69 → 140
# ═══════════════════════════════════════════════════════════════════════════
print(f"\n\n{'='*70}")
print("CUMULATIVE SAMPLE BUILD-UP")
print("=" * 70)

eba_set = set(EBA_COUNTRIES)
ssa_set = set(SSA_COUNTRIES)
orig_69 = eba_set | ssa_set

samples_fixed = [
    ('EBA-49', eba_set),
    ('Original-69', orig_69),
    ('Full-140', set(all_countries)),
]

print(f"\n  {'Sample':<15s} {'N_c':>4s} {'N_obs':>6s} {'R²':>7s} {'Z₁':>9s} {'p(Z₁)':>8s} {'Z₂':>9s} {'p(Z₂)':>8s} {'Z₃':>9s} {'p(Z₃)':>8s}")
print("  " + "-" * 88)

for name, country_set in samples_fixed:
    sub = base_sample[base_sample['iso3'].isin(country_set)].copy()
    n_c = sub['iso3'].nunique()
    r = run_gls(sub, 'ca_gdp', base_vars)
    z1c, z1p = r['coefficients']['Z_1'], r['p_values']['Z_1']
    z2c, z2p = r['coefficients']['Z_2'], r['p_values']['Z_2']
    z3c, z3p = r['coefficients']['Z_3'], r['p_values']['Z_3']
    print(f"  {name:<15s} {n_c:4d} {r['n_obs']:6d} {r['r_squared']:7.4f} {z1c:9.2f} {z1p:8.4f} {z2c:9.3f} {z2p:8.4f} {z3c:9.4f} {z3p:8.4f}")

# Cumulative region addition starting from Original-69
print(f"\n--- Cumulative Region Addition (starting from Original-69) ---")
region_add_order = ['Adv_Europe', 'E_Asia', 'Adv_Other', 'LatAm', 'MENA', 'S/SE_Asia', 'CCA', 'Other']

expansion_regions = []
for reg in region_add_order:
    new_countries = set(base_sample[base_sample['region'] == reg]['iso3'].unique()) - orig_69
    if new_countries:
        expansion_regions.append((reg, new_countries))

included = set(orig_69) & set(base_sample['iso3'].unique())

print(f"\n  {'Cumulative':<35s} {'N_c':>4s} {'N_obs':>6s} {'R²':>7s} {'Z₁':>9s} {'p(Z₁)':>8s} {'Z₂':>9s} {'p(Z₂)':>8s} {'Z₃':>9s} {'p(Z₃)':>8s}")
print("  " + "-" * 108)

sub = base_sample[base_sample['iso3'].isin(included)].copy()
r = run_gls(sub, 'ca_gdp', base_vars)
z1c, z1p = r['coefficients']['Z_1'], r['p_values']['Z_1']
z2c, z2p = r['coefficients']['Z_2'], r['p_values']['Z_2']
z3c, z3p = r['coefficients']['Z_3'], r['p_values']['Z_3']
label = f"Original-69 ({sub['iso3'].nunique()})"
print(f"  {label:<33s} {sub['iso3'].nunique():4d} {r['n_obs']:6d} {r['r_squared']:7.4f} {z1c:9.2f} {z1p:8.4f} {z2c:9.3f} {z2p:8.4f} {z3c:9.4f} {z3p:8.4f}")

for reg, new_cs in expansion_regions:
    included = included | new_cs
    sub = base_sample[base_sample['iso3'].isin(included)].copy()
    r = run_gls(sub, 'ca_gdp', base_vars)
    z1c, z1p = r['coefficients']['Z_1'], r['p_values']['Z_1']
    z2c, z2p = r['coefficients']['Z_2'], r['p_values']['Z_2']
    z3c, z3p = r['coefficients']['Z_3'], r['p_values']['Z_3']
    label = f"+ {reg} (+{len(new_cs)})"
    print(f"  {label:<33s} {sub['iso3'].nunique():4d} {r['n_obs']:6d} {r['r_squared']:7.4f} {z1c:9.2f} {z1p:8.4f} {z2c:9.3f} {z2p:8.4f} {z3c:9.4f} {z3p:8.4f}")


# Save all results
jk_df.to_csv(OUTPUT_DIR / 'jackknife_baseline.csv', index=False)
if ext_jk_results:
    ext_jk_df.to_csv(OUTPUT_DIR / 'jackknife_extended.csv', index=False)
print(f"\nSaved to {OUTPUT_DIR / 'jackknife_baseline.csv'}")
print("Done.")
